#!/usr/bin/env python3
# -*- coding: utf-8 -*-

''' 
    date: 2021-03-19
    author: baole 

    模块功能：
        从东财爬取大股东所持有的股本数量,并计算实际流通股数
        原网址：http://f10.eastmoney.com/ShareholderResearch/Index?type=web&code=SZ300171
        里面的单位是股。

        遍历所有的股票代码，获取大股东持股，计算实际流通股数，保存到数据库中。
'''

import os
import sys
import requests
import json
import pymysql
import time
import random


from quantway.lib.utils import read_jsonparam, get_module_name, get_codelist, split_list_by_num



def get_dagudong(confpath, pid, li_stock_code):
    '''
        从东财爬取大股东持股数
    '''

    # 获取参数
    param = read_jsonparam(confpath)
    db_param = param.database

    try:
        db = pymysql.connect(**db_param)
    except Exception as e:
        print("[PID=%d][%s]连接数据库失败,程序终止" % (pid, get_module_name()), e)
        return
    else:
        print("[PID=%d][%s]:数据库连接成功" % (pid, get_module_name()))


    cursor = db.cursor()
    li_fail = list()                    # 失败的股票列表
    num_total = len(li_stock_code)      # 本进程处理的股票总数
    num_fail = 0                        # 失败数量

    for i, stock_code in enumerate(li_stock_code):
        # 随机休息1~20s， 否则容易被东财屏蔽ip一段时间
        # time.sleep(random.randint(1,20))
        
        # 抓取数据
        r = query_gudong_from_eastmoney(stock_code.upper())
        
        if r is None:
            num_fail += 1 
            li_fail.append(stock_code)
            print("[PID=%d][%d/%d] 获取%s的数据失败" % (pid, i+1, num_total, stock_code))
            continue
        shareholder = json.loads(r.text)
        if shareholder == None or len(shareholder) == 0:
            num_fail += 1 
            li_fail.append(stock_code)
            print("[PID=%d][%d/%d] 获取%s的数据失败" % (pid, i+1, num_total, stock_code))
            continue 
        if 'status' in shareholder and shareholder['status'] == -1:
            num_fail += 1 
            li_fail.append(stock_code)
            print("[PID=%d][%d/%d] 获取%s的数据失败 %s" % (pid, i+1, num_total, stock_code, shareholder['message']))
            continue

        # >5%的股东总持股数（单位：股）
        lt5cgs = 0
        try:
            # 总股本
            if len(shareholder['sdltgd']) > 0:
                last_shareholder = shareholder['sdltgd'][0]
                li_shareholder = last_shareholder['sdltgd']
                for shareholder in li_shareholder:
                    # 获取股份类型，判断是否是A股
                    gflx = shareholder['gflx']  
                    if not "A" in gflx:
                        continue
                    # 获取“占流通股持股比例"
                    zltgbcgbl = shareholder['zltgbcgbl']
                    if float(zltgbcgbl[:-1]) >= 5:
                        # 获取"持股数"(单位：股)
                        cgs = float(shareholder['cgs'].replace(",",""))
                        lt5cgs += cgs
                
        except Exception as e:
            num_fail += 1  
            li_fail.append(stock_code)
            print("[PID=%d][%d/%d] 获取%s的数据失败" % (pid, i+1, num_total, stock_code), e)
            continue
        
        try:
            db.begin()
            # 读取这个股票的以上市流通A股数量（万股）
            sql = "select stock_code, yssltag from guben where stock_code='%s'" % stock_code
            if cursor.execute(sql) > 0:
                row = cursor.fetchone()
                yssltag = row[1]
                # 计算实际流通A股（单位：股）
                sjltag =  yssltag * 10000 - lt5cgs
                # 更新总股本和流通A股
                sql = "update guben set lt5cgs=%.2f ,sjltag=%.2f where stock_code='%s'" % (lt5cgs, sjltag, stock_code)
                cursor.execute(sql)
                db.commit()
                print("[PID=%d][%d/%d] 股票%s的流通A股%.2f, 大股东持有数%.2f, 实际流通A股%.2f " % (pid, i+1, num_total, stock_code, yssltag*10000, lt5cgs, sjltag))
            else:
                num_fail += 1 
                li_fail.append(stock_code)
                print("[PID=%d][%d/%d] 数据库中没有%s的股本和流通A股" % (pid, i+1, num_total, stock_code ))


        except Exception as e:
            db.rollback()
            num_fail += 1 
            li_fail.append(stock_code)
            print("[PID=%d][%d/%d]%s保存数据失败" % (pid, i+1, num_total, stock_code), e)


    cursor.close()
    db.close()

    print("[PID=%d]总共:%d股票，失败%d股票" % (pid, num_total, num_fail))
    if num_fail > 0:
        print("[PID=%d]失败的股票：" % pid , ",".join(li_fail))



def query_gudong_from_eastmoney(stock_code):
    # 对于失败的request,我们会尝试5次。
    attempt = 0
    while attempt<5:
        try:
            r = requests.get('http://f10.eastmoney.com/ShareholderResearch/ShareholderResearchAjax?code=%s' % stock_code)
            return r                
        except Exception as e:
            print(e)
            time.sleep(60)
        attempt+=1
        print("获取%s的股本超时，重试%d " % (stock_code, attempt) )


    return None


        


    


if __name__ == '__main__':

    # 去掉科创板
    li_stock_code = [code for code in get_codelist("all_stockcode") if not code.startswith("sh68")]

    if len(sys.argv) > 2 and ("--single" not in sys.argv):
        total = int(sys.argv[1])
        pid = int(sys.argv[2])
        split = split_list_by_num(li_stock_code, total)
        get_dagudong("conf/spider.json", pid, split[pid])
    else:
        if "--single" in sys.argv:
            li_stock_code = [sys.argv[sys.argv.index("--single")+1]]
        get_dagudong("conf/spider.json", 0, li_stock_code)

    
